import itertools

import torch
from torch.distributed._tensor import distribute_tensor, Replicate, Shard
from torch.testing._internal.common_utils import (
    instantiate_parametrized_tests,
    parametrize,
    run_tests,
)

import torch_npu
from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase
from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU


def get_shape_from_layout(batch: int, num_head: int, seq_length: int, dimension: int, layout: str):
    layout_map = {
        "B": batch,
        "N": num_head,
        "S": seq_length,
        "D": dimension,
        "1": 1,
    }
    shape = []
    for dim in layout:
        if dim in layout_map:
            shape.append(layout_map[dim])
        else:
            raise ValueError(f"Invalid layout character: {dim}")

    return tuple(shape)


class TestMathOps(NPUDTensorTestBase):
    @skipIfUnsupportMultiNPU(4)
    @with_comms
    def test_npu_rms_norm_forward(self):
        device_mesh = self.build_device_mesh()

        x = torch.randn((1, 128, 64), dtype=torch.float32).npu()
        gamma = torch.randn(64, dtype=torch.float32).npu()

        y, rstd = torch_npu.npu_rms_norm(x, gamma)

        dist_x = distribute_tensor(x, device_mesh, [Shard(1)])
        dist_gamma = distribute_tensor(gamma, device_mesh, [Replicate()])

        dist_y, dist_rstd = torch_npu.npu_rms_norm(dist_x, dist_gamma)

        self.assertEqual(dist_y.full_tensor(), y)
        self.assertEqual(dist_gamma.full_tensor(), gamma)

    @skipIfUnsupportMultiNPU(4)
    @with_comms
    def test_npu_rms_norm_backward(self):
        device_mesh = self.build_device_mesh()

        x = torch.randn((1, 128, 64), dtype=torch.float32).npu()
        gamma = torch.randn(64, dtype=torch.float32).npu()
        grad_y = torch.randn((1, 128, 64), dtype=torch.float32).npu()

        x = x.npu()
        gamma = gamma.npu()
        grad_y = grad_y.npu()
        x.requires_grad = True
        gamma.requires_grad = True

        y, rstd = torch_npu.npu_rms_norm(x, gamma, epsilon=1e-06)
        y.backward(grad_y)
        dx = x.grad
        dw = gamma.grad

        dist_x = distribute_tensor(x, device_mesh, [Shard(2)])
        dist_gamma = distribute_tensor(gamma, device_mesh, [Replicate()])

        dist_y, dist_rsts = torch_npu.npu_rms_norm(dist_x, dist_gamma, epsilon=1e-06)
        dist_grad_y = distribute_tensor(grad_y, device_mesh, dist_y.placements)
        dist_y.backward(dist_grad_y)
        dist_dx = dist_x.grad
        dist_dw = dist_gamma.grad

        self.assertEqual(dist_y.full_tensor(), y)
        self.assertEqual(dist_gamma.full_tensor(), gamma)

        self.assertEqual(dist_dx.full_tensor(), dx)
        self.assertEqual(dist_dw.full_tensor(), dw)

    @skipIfUnsupportMultiNPU(4)
    @with_comms
    def test_npu_add_rms_norm_forward(self):
        device_mesh = self.build_device_mesh()

        x1 = torch.randn((1, 128, 64), dtype=torch.float32).npu()
        x2 = torch.randn((1, 128, 64), dtype=torch.float32).npu()
        gamma = torch.randn(64, dtype=torch.float32).npu()

        y, rstd, x = torch_npu.npu_add_rms_norm(x1, x2, gamma)

        def test_placement_comb(placements1, placements2):
            dist_x1 = distribute_tensor(x1, device_mesh, placements1)
            dist_x2 = distribute_tensor(x2, device_mesh, placements2)
            dist_gamma = distribute_tensor(gamma, device_mesh, [Replicate()])
            dist_y, dist_rstd, dist_x = torch_npu.npu_add_rms_norm(dist_x1, dist_x2, dist_gamma)
            self.assertEqual(dist_y.full_tensor(), y)
            self.assertEqual(dist_rstd.full_tensor(), rstd)
            self.assertEqual(dist_x.full_tensor(), x)

        placement = [Shard(0), Shard(1), Shard(2), Replicate()]
        placement_combs = itertools.product(placement, placement)
        for comb in placement_combs:
            test_placement_comb([comb[0]], [comb[1]])

    @skipIfUnsupportMultiNPU(4)
    @with_comms
    @parametrize(
        "rotary_mode,input_layout,sin_cos_layout",
        [
            ("half", "BNSD", "11SD"),
            ("half", "BNSD", "B1SD"),
            ("half", "BNSD", "BNSD"),
            ("half", "BSND", "1S1D"),
            ("half", "BSND", "BS1D"),
            ("half", "BSND", "BSND"),
            ("half", "SBND", "S11D"),
            ("half", "SBND", "SB1D"),
            ("half", "SBND", "SBND"),
            ("interleave", "BNSD", "11SD"),
            ("interleave", "BSND", "1S1D"),
            ("interleave", "SBND", "S11D"),
        ]
    )
    def test_npu_rotary_mul_forward(self, rotary_mode, input_layout, sin_cos_layout):
        device_mesh = self.build_device_mesh()

        B = 8
        N = 8
        S = 64
        D = 32
        x_shape = get_shape_from_layout(B, N, S, D, input_layout)
        x = torch.randn(x_shape, dtype=torch.float32, device="npu")
        sin_cos_shape = get_shape_from_layout(B, N, S, D, sin_cos_layout)
        sin = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
        cos = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1

        y = torch_npu.npu_rotary_mul(x, cos, sin, rotary_mode=rotary_mode)

        def test_placement_comb(x_placements, sin_placements, cos_placements):
            dist_x = distribute_tensor(x, device_mesh, x_placements)
            dist_sin = distribute_tensor(sin, device_mesh, sin_placements)
            dist_cos = distribute_tensor(cos, device_mesh, cos_placements)
            dist_y = torch_npu.npu_rotary_mul(dist_x, dist_cos, dist_sin, rotary_mode=rotary_mode)
            self.assertEqual(dist_y.full_tensor(), y)

        placements = [Shard(0), Shard(1), Shard(2), Replicate()]
        for placement in placements:
            if isinstance(placement, Shard) and sin_cos_shape[placement.dim] == 1:
                test_placement_comb([placement], [Replicate()], [Replicate()])
            else:
                test_placement_comb([placement], [placement], [placement])

    @skipIfUnsupportMultiNPU(4)
    @with_comms
    @parametrize(
        "rotary_mode,input_layout,sin_cos_layout",
        [
            ("half", "BNSD", "11SD"),
            ("half", "BNSD", "B1SD"),
            ("half", "BNSD", "BNSD"),
            ("half", "BSND", "1S1D"),
            ("half", "BSND", "BS1D"),
            ("half", "BSND", "BSND"),
            ("half", "SBND", "S11D"),
            ("half", "SBND", "SB1D"),
            ("half", "SBND", "SBND"),
            ("interleave", "BNSD", "11SD"),
            ("interleave", "BSND", "1S1D"),
            ("interleave", "SBND", "S11D"),
        ]
    )
    def test_npu_rotary_mul_backward(self, rotary_mode, input_layout, sin_cos_layout):
        device_mesh = self.build_device_mesh()

        B = 8
        N = 8
        S = 64
        D = 32
        x_shape = get_shape_from_layout(B, N, S, D, input_layout)
        x = torch.randn(x_shape, dtype=torch.float32, device="npu", requires_grad=True)
        sin_cos_shape = get_shape_from_layout(B, N, S, D, sin_cos_layout)
        sin = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
        cos = torch.randn(sin_cos_shape, dtype=torch.float32, device="npu") * 2 - 1
        sin.requires_grad = True
        cos.requires_grad = True


        y = torch_npu.npu_rotary_mul(x, cos, sin, rotary_mode=rotary_mode)
        grad_y = torch.ones_like(y, dtype=torch.float32, device="npu")
        y.backward(grad_y)

        def test_placement_comb(x_placements, sin_placements, cos_placements):
            dist_x = distribute_tensor(x, device_mesh, x_placements)
            dist_sin = distribute_tensor(sin, device_mesh, sin_placements)
            dist_cos = distribute_tensor(cos, device_mesh, cos_placements)
            dist_y = torch_npu.npu_rotary_mul(dist_x, dist_cos, dist_sin, rotary_mode=rotary_mode)
            dist_grad_y = distribute_tensor(grad_y, device_mesh, dist_y.placements)
            dist_y.backward(dist_grad_y)
            self.assertEqual(dist_y.full_tensor(), y)
            self.assertEqual(dist_x.grad.full_tensor(), x.grad)
            self.assertEqual(dist_sin.grad.full_tensor(), sin.grad)
            self.assertEqual(dist_cos.grad.full_tensor(), cos.grad)

        placements = [Shard(0), Shard(1), Shard(2), Replicate()]
        for placement in placements:
            if isinstance(placement, Shard) and sin_cos_shape[placement.dim] == 1:
                test_placement_comb([placement], [Replicate()], [Replicate()])
            else:
                test_placement_comb([placement], [placement], [placement])


instantiate_parametrized_tests(TestMathOps)


if __name__ == "__main__":
    run_tests()
